#include "pmm.h"
#include "printf.h"
#include "vmm.h"  // 引入PGSIZE等宏定义，确保页大小一致
#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>

// 物理内存管理器状态
static struct {
    uintptr_t memory_start;  // 可用内存起始地址（内核结束后，位图之后）
    uintptr_t memory_end;    // 可用内存结束地址（物理内存上限）
    size_t total_pages;      // 总页数
    size_t free_pages;       // 空闲页数
    uint8_t* page_bitmap;    // 页位图: 1=已分配, 0=空闲
    uintptr_t physical_start; // 物理内存起始地址（传入的start，按页对齐后）
} pmm;

// 计算页索引（相对于物理内存起始地址）
static size_t page_index(void* page) {
    uintptr_t addr = (uintptr_t)page;
    // 确保地址是页对齐的（PGSIZE=4096）
    if (addr % PGSIZE != 0) {
        return (size_t)-1;
    }
    
    // 计算相对于物理内存起始地址的偏移量对应的页索引
    return (addr - pmm.physical_start) / PGSIZE;
}

// 初始化物理内存管理器（修改为带参数：物理内存起始地址和结束地址）
void pmm_init(uint64_t start, uint64_t end) {
    // 1. 对传入的起始地址（内核结束地址）按页对齐
    uintptr_t kernel_end = start;
    uintptr_t kernel_end_aligned = (kernel_end + PGSIZE - 1) & ~(PGSIZE - 1);
    pmm.physical_start = kernel_end_aligned;  // 物理内存管理的起始地址（对齐后）
    
    // 2. 物理内存结束地址（不超过传入的end）
    pmm.memory_end = end;
    if (pmm.memory_end <= pmm.physical_start) {
        panic("pmm_init: 物理内存结束地址小于起始地址（%p <= %p）", 
              (void*)pmm.memory_end, (void*)pmm.physical_start);
    }
    
    // 3. 计算总页数
    size_t total_memory = pmm.memory_end - pmm.physical_start;
    pmm.total_pages = total_memory / PGSIZE;
    if (pmm.total_pages == 0) {
        panic("pmm_init: 可用物理内存不足一页（%d bytes）", total_memory);
    }
    
    // 4. 计算页位图所需大小并分配（位图存储在物理内存最开始位置）
    size_t bitmap_size = (pmm.total_pages + 7) / 8;  // 向上取整（1位表示1页）
    pmm.page_bitmap = (uint8_t*)pmm.physical_start;  // 位图起始地址
    
    // 5. 计算位图占用的页数，并标记为已分配
    size_t bitmap_pages = (bitmap_size + PGSIZE - 1) / PGSIZE;
    pmm.memory_start = pmm.physical_start + bitmap_pages * PGSIZE;  // 可用内存起始（位图之后）
    
    // 检查位图是否超出物理内存范围
    if (pmm.memory_start >= pmm.memory_end) {
        panic("pmm_init: 页位图占用内存超出物理内存范围（%p >= %p）",
              (void*)pmm.memory_start, (void*)pmm.memory_end);
    }
    
    // 6. 初始化空闲页数（总页数 - 位图占用页数）
    pmm.free_pages = pmm.total_pages - bitmap_pages;
    
    // 7. 初始化位图：先全部置为1（已分配）
    for (size_t i = 0; i < bitmap_size; i++) {
        pmm.page_bitmap[i] = 0xff;
    }
    
    // 8. 标记可用内存页为0（空闲）
    size_t first_free_page = (pmm.memory_start - pmm.physical_start) / PGSIZE;
    for (size_t i = first_free_page; i < pmm.total_pages; i++) {
        pmm.page_bitmap[i / 8] &= ~(1 << (i % 8));
    }
    
    // 输出初始化信息（验证内存范围）
    printf("物理内存管理器初始化完成:\n");
    printf("  内核结束地址: 0x%lx (对齐后: 0x%lx)\n", kernel_end, kernel_end_aligned);
    printf("  内存范围: 0x%lx - 0x%lx\n", pmm.memory_start, pmm.memory_end);
    printf("  总页数: %d\n", pmm.total_pages);
    printf("  空闲页数: %d\n", pmm.free_pages);
    printf("  页大小: %d bytes\n", PGSIZE);
    printf("  位图大小: %d bytes (占用 %d 页)\n", bitmap_size, bitmap_pages);
}

// 分配一个物理页
void* alloc_page(void) {
    return alloc_pages(1);
}

// 分配n个连续的物理页
void* alloc_pages(size_t n) {
    if (n == 0 || n > pmm.free_pages) {
        return NULL;
    }
    
    size_t consecutive = 0;
    size_t start_index = 0;
    
    // 查找连续的n个空闲页
    for (size_t i = 0; i < pmm.total_pages; i++) {
        // 检查第i位是否为0（空闲）
        if (!(pmm.page_bitmap[i / 8] & (1 << (i % 8)))) {
            consecutive++;
            if (consecutive == n) {
                start_index = i - n + 1;
                break;
            }
        } else {
            consecutive = 0;
        }
    }
    
    // 未找到足够的连续空闲页
    if (consecutive < n) {
        return NULL;
    }
    
    // 标记这些页为已分配
    for (size_t i = start_index; i < start_index + n; i++) {
        pmm.page_bitmap[i / 8] |= (1 << (i % 8));
    }
    
    pmm.free_pages -= n;
    
    // 计算物理页地址并返回
    return (void*)(pmm.physical_start + start_index * PGSIZE);
}

// 释放一个物理页
void free_page(void* page) {
    free_pages(page, 1);
}

// 释放n个连续的物理页
void free_pages(void* page, size_t n) {
    if (page == NULL || n == 0) {
        return;
    }
    
    uintptr_t addr = (uintptr_t)page;
    // 检查地址是否在管理范围内且页对齐
    if (addr < pmm.memory_start || addr >= pmm.memory_end || addr % PGSIZE != 0) {
        printf("[PMM] 释放失败：地址无效 0x%lx\n", addr);
        return;
    }
    
    size_t start_index = page_index(page);
    // 检查索引有效性
    if (start_index == (size_t)-1 || start_index + n > pmm.total_pages) {
        printf("[PMM] 释放失败：索引越界 %d\n", start_index);
        return;
    }
    
    // 检查这些页是否都已分配
    for (size_t i = start_index; i < start_index + n; i++) {
        if (!(pmm.page_bitmap[i / 8] & (1 << (i % 8)))) {
            printf("[PMM] 释放失败：页未分配 %d\n", i);
            return;
        }
    }
    
    // 标记这些页为空闲
    for (size_t i = start_index; i < start_index + n; i++) {
        pmm.page_bitmap[i / 8] &= ~(1 << (i % 8));
    }
    
    pmm.free_pages += n;
}

// 获取内存信息
void get_memory_info(size_t* total, size_t* free) {
    if (total) {
        *total = pmm.total_pages * PGSIZE;
    }
    if (free) {
        *free = pmm.free_pages * PGSIZE;
    }
}
